package org.apache.lorry.zookeeper.meeting;

import lombok.extern.slf4j.Slf4j;
import org.apache.zookeeper.CreateMode;
import org.apache.zookeeper.KeeperException;
import org.apache.zookeeper.ZooDefs;
import org.apache.zookeeper.ZooKeeper;

import java.util.ArrayList;
import java.util.List;

@Slf4j
public class RecoveredAssignments {
    List<String> tasks;

    List<String> assignments;   //已经分配了的任务

    List<String> status;    //已经完成了的任务

    List<String> activeWorkers;

    List<String> assignedWorkers;

    RecoveryCallback cb;

    ZooKeeper zooKeeper;


    public interface RecoveryCallback {
        final static int OK = 0;

        final static int FAILED = -1;

        public void recoveryComplete(int rc, List<String> tasks);
    }


    public RecoveredAssignments(ZooKeeper zooKeeper) {
        this.zooKeeper = zooKeeper;
        this.assignments = new ArrayList<>();
    }

    public void recover(RecoveryCallback cb) {
        this.cb = cb;
        getTasks();
    }

    private void getTasks() {
        zooKeeper.getChildren(ZkConstant.TASK_PATH, false, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getTasks();
                    break;
                case OK:
                    log.info("Getting tasks for recovery");
                    tasks = children;
                    getAssignedWorkers();
                    break;
                default:
                    log.error("getChildren failed",  KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null);
            }
        }, null);
    }


    private void getAssignedWorkers() {
        zooKeeper.getChildren(ZkConstant.ASSIGN_PATH, false, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getAssignedWorkers();
                    break;
                case OK:
                    assignedWorkers = children;
                    getWorkers(children);
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null);
            }
        }, null);
    }


    private void getWorkers(Object ctx) {
        zooKeeper.getChildren(ZkConstant.LISTENER_PATH, false, (rc, path, ctx1, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getWorkers(ctx1);
                    break;
                case OK:
                    log.info("Getting worker assignments for recovery: " + children.size());

                    if (children.size() == 0) {
                        log.warn("Empty list of workers, possibly just starting");
                        cb.recoveryComplete(RecoveryCallback.OK, new ArrayList<>());
                        break;
                    }
                    /*
                     * Need to know which of the assigned workers are active.
                     */
                    activeWorkers = children;
                    for (String s : assignedWorkers) {
                        getWorkerAssignments(ZkConstant.ASSIGN_PATH + "/" + s);
                    }
                    break;
                default:
                    log.error("getChildren failed",  KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(RecoveryCallback.FAILED, null);
            }
        }, ctx);
    }


    private void getWorkerAssignments(String path) {
        zooKeeper.getChildren(path, false,
                (rc, p, ctx, children) -> {
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            getWorkerAssignments(path);
                            break;
                        case OK:
                            String worker = path.replace(ZkConstant.ASSIGN_PATH, "");
                            /*
                             * Reassign the tasks of this worker.
                             */
                            if(activeWorkers.contains(worker)) {
                                assignments.addAll(children);   //已经分配了的任务
                            }else{
                                for(String task : children){
                                    if(!tasks.contains(task)){
                                        tasks.add(task);
                                        getDataReassign(path, task);
                                    }else{
                                        deleteAssignment(path + "/" + task);
                                    }
                                    deleteAssignment(path);
                                }
                            }
                            assignedWorkers.remove(worker);

                            if(assignedWorkers.size() == 0){
                                getStatuses();
                            }
                            break;
                        case NONODE:
                            log.info("No such znode exists : {}", path);
                            break;
                        default:
                            log.error("getChildren failed",  KeeperException.create(KeeperException.Code.get(rc), path));
                            cb.recoveryComplete(RecoveryCallback.FAILED, null);
                    }
                }, null);
    }


    void getStatuses() {
        zooKeeper.getChildren(ZkConstant.STATUS_PATH, false, (rc, path, ctx, children) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getStatuses();
                    break;
                case OK:
                    log.info("Processing assignments for recovery");
                    status = children;
                    processAssignments();
                    break;
                default:
                    log.error("getChildren failed", KeeperException.create(KeeperException.Code.get(rc), path));
                    cb.recoveryComplete(org.apache.lorry.zookeeper.book.recovery.RecoveredAssignments.RecoveryCallback.FAILED, null);
            }
        }, null);
    }

    private void processAssignments(){
        log.info("Size of tasks: " + tasks.size());
        for (String s : assignments) {
            log.info("Assignment: " + s);
            deleteAssignment(ZkConstant.TASK_PATH + "/" + s);
            tasks.remove(s);
        }
        log.info("Size of tasks after assignment filtering: " + tasks.size());
        for(String s: status){
            log.info("Checking task: {} ", s);
            deleteAssignment(ZkConstant.TASK_PATH + "/" + s);
            tasks.remove(s);
        }
        log.info("Size of tasks after status filtering: " + tasks.size());

        cb.recoveryComplete(RecoveryCallback.OK, tasks);
    }

    private void getDataReassign(String path, String task) {
        zooKeeper.getData(path, false, (rc, p, ctx, data, stat) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    getDataReassign(path, (String) ctx);
                    break;
                case OK:
                    recreateTask(new RecreateTaskCtx(path, task, data));
                    break;
                default:
                    log.error("Something went wrong when getting data ",
                            KeeperException.create(KeeperException.Code.get(rc)));
            }
        }, null);
    }

    void recreateTask(RecreateTaskCtx ctx) {
        zooKeeper.create(ZkConstant.TASK_PATH + "/" + ctx.task, ctx.data, ZooDefs.Ids.OPEN_ACL_UNSAFE, CreateMode.PERSISTENT,
                (int rc, String path, Object ctx1, String name) ->{
                    switch (KeeperException.Code.get(rc)) {
                        case CONNECTIONLOSS:
                            recreateTask((RecreateTaskCtx) ctx1);
                            break;
                        case OK:
                            deleteAssignment(((RecreateTaskCtx) ctx1).path);
                            break;
                        case NODEEXISTS:
                            log.warn("Node shouldn't exist: " + path);
                            break;
                        default:
                            log.error("Something wwnt wrong when recreating task", KeeperException.create(KeeperException.Code.get(rc), path));
                    }
                }, ctx);
    }


    void deleteAssignment(String path) {
        zooKeeper.delete(path, -1, (int rc, String p, Object ctx) -> {
            switch (KeeperException.Code.get(rc)) {
                case CONNECTIONLOSS:
                    deleteAssignment(path);
                    break;
                case OK:
                    log.info("Task correctly deleted: " + path);
                    break;
                default:
                    log.error("Failed to delete task data", KeeperException.create(KeeperException.Code.get(rc), path));
                    break;
            }
        }, null);
    }

    class RecreateTaskCtx {
        String path;
        String task;
        byte[] data;

        RecreateTaskCtx(String path, String task, byte[] data) {
            this.path = path;
            this.task = task;
            this.data = data;
        }
    }


}